from vec import Vec2
import math
import random
import pgzrun

WIDTH = 500
HEIGHT = 500

PEN_COLOR = (0, 200, 0)
BACK_COLOR = (0, 0, 0)
ptList = []
length = 300
pt_mid = Vec2(WIDTH / 2, HEIGHT / 2)
y_offset = length / math.sqrt(3) 
pt_a = Vec2(pt_mid.x, pt_mid.y - y_offset)
pt_d = Vec2(pt_mid.x, pt_mid.y + length / 2 / math.sqrt(3))
pt_b = Vec2(pt_d.x - length / 2, pt_d.y)
pt_c = Vec2(pt_d.x + length / 2, pt_d.y)

ptList.append(pt_a)
ptList.append(pt_b)
ptList.append(pt_c)

def distort(p):
    r = random.randint(1, 3)
    vx = random.randint(-100, 100)
    vy = random.randint(-100, 100)
    v = Vec2(vx, vy)
    v.normalize()
    dp = p + v * r

    return dp

def draw_fractal(pa, pb, n):
    if n == 1:
        if enableDistort == False:
            screen.draw.line(pa.raw_get(), pb.raw_get(), PEN_COLOR)
        else:
            screen.draw.line(distort(pa).raw_get(), distort(pb).raw_get(), PEN_COLOR)
    else:
        pts = []
        dir = pb - pa
        edgeLen = dir.length()
        dir.normalize()
        midPt = (pa + pb) / 2
        pts.append(pa)

        pt1 = pa * 2 / 3 + pb / 3

        pts.append(pt1)

        vdir = Vec2(0, 0)
        vdir.x = -dir.y
        vdir.y = dir.x

        dist = edgeLen / 3 / 2 * math.sqrt(3)

        pt2 = midPt + vdir * dist

        pts.append(pt2)

        pt3 = pa / 3 + pb * 2 / 3

        pts.append(pt3)

        pts.append(pb)

        for i in range(len(pts) - 1):
            draw_fractal(pts[i], pts[i + 1], n - 1)

count = len(ptList)
n = 5
enableDistort = False

def update():
    pass

def draw():
    screen.fill(BACK_COLOR)
    
    for i in range(count):
        if i == count - 1:
            draw_fractal(ptList[count - 1], ptList[0], n)
        else:
            draw_fractal(ptList[i], ptList[i + 1], n) 

pgzrun.go()